
import sys
from coverage import coveragefromsummary
import indexing.featureindex

# KNOWN_GENE_TABLE = "/home/bowtie/Bowtie_Illumina_Alignment/references/mergedGeneCluster-plus-additional.table.txt"
KNOWN_GENE_TABLE = "SureSelect-genes.txt"

def findGenes(query, filename=KNOWN_GENE_TABLE):
  (chr, strand, spos, epos) = query
  genes = []
  for line in open(filename):
    if line[0]=="#": continue
    cols = line.strip().split("\t")
    if chr!=cols[1]: continue
    if strand!=cols[2]: continue
    if long(cols[4])>=spos and long(cols[3])<=epos: genes.append(cols[0])
  return genes

def aggregateGeneCoverage(genes, summaryFiles, filename=KNOWN_GENE_TABLE, outputstream=sys.stdout, humanMappedReads=None):
  if genes==None or len(genes)==0:
    return humanMappedReads
  
  exons = {}  
  introns = {}
  chr = None
  spos = -1
  epos = -1
  notFound = {}
  geneStrand = {}
  for gene in genes: notFound[gene] = True
  for line in open(filename):
    if line[0]=="#": continue
    cols = line.strip().split("\t")
    if not cols[0] in genes: continue
    gene = cols[0]
    if chr==None: chr = cols[1]
    elif chr!=cols[1]: raise Exception("bad chr")
    geneStrand[gene] = cols[2]
    if spos<0 or spos>long(cols[3]): spos = long(cols[3])
    if epos<0 or epos<long(cols[4]): epos = long(cols[4])
    notFound[gene] = False
    exonCount = int(cols[7])
    exonStarts = cols[8].split(",")[:-1]
    exonEnds = cols[9].split(",")[:-1]
    exons[gene] = []
    introns[gene] = []
    for i in xrange(exonCount):
      exons[gene].append((long(exonStarts[i]), long(exonEnds[i]), i+1))
    for i in xrange(exonCount-1):
      if exons[gene][i+1][0] > exons[gene][i][1]+1:
        introns[gene].append((exons[gene][i][1]+1, exons[gene][i+1][0]-1, i+1))
  for gene in genes:
    if notFound[gene] == True: raise Exception("gene not found: " + gene)
  query = (chr, "?", spos, epos)
  print >> sys.stderr, "# Query =", query

  if humanMappedReads!=None and isinstance(humanMappedReads, list):
    print >> sys.stderr, "# Counting human-mapped reads...."
    for summaryFile in summaryFiles:
      numReads = 0
      ix = indexing.featureindex.FeatureIndex(summaryFile)
      for feature in ix.getFeatures():
        details = feature.getDetails()
        if details.startDrop == "0":
          numReads += 1
      humanMappedReads.append(numReads)
      print >> sys.stderr, "# Human-mapped reads in", summaryFile, "=", numReads
    print >> sys.stderr, "# Final counts: ", humanMappedReads

  coverage = [0]*(epos-spos+1)
  for summaryFile in summaryFiles:
    for (x, cts) in coveragefromsummary.process(query, summaryFile, coveragefromsummary.SkipZeroes.ALL):
      coverage[x-spos] += cts.total

  qs = query[2]
  for gene in genes:
    strand = geneStrand[gene]
    geneName = gene.replace(" ", "_")
    for exon in exons[gene]:
      cov = coverage[exon[0]-qs:exon[1]-qs+1]
      print >> outputstream, "%s\tExon\t%d\t%s\t%s\t%s\t%ld\t%ld\t%f\t%f" % (geneName, exon[2], "%s_%s_%s_%s" % (chr, strand, exon[0], exon[1]), chr, strand, exon[0], exon[1], sum(cov), sum([x*x for x in cov]))
    for intron in introns[gene]:
      cov = coverage[intron[0]-qs:intron[1]-qs+1]
      print >> outputstream, "%s\tIntron\t%d\t%s\t%s\t%s\t%ld\t%ld\t%f\t%f" % (geneName, intron[2], "%s_%s_%s_%s" % (chr, strand, intron[0], intron[1]), chr, strand, intron[0], intron[1], sum(cov), sum([x*x for x in cov]))

  return humanMappedReads

def aggregateAllChromosomes(summaryFiles, geneFile=KNOWN_GENE_TABLE, outputstream=sys.stdout):
  humanMappedReads = []
  for chr in ["X", "Y"] + range(1, 23):
    query1 = ("chr" + str(chr), "+", 0, 1000000000L)
    query2 = ("chr" + str(chr), "-", 0, 1000000000L)
    genes = findGenes(query1, geneFile)
    genes.extend(findGenes(query2, geneFile))
    print >> sys.stderr, "# Genes in", query1, "and", query2, "=", len(genes)
    mappedReads = None
    if humanMappedReads==[] and len(genes)>0:
      mappedReads = humanMappedReads
    aggregateGeneCoverage(genes, summaryFiles, geneFile, outputstream, mappedReads)
    if mappedReads != None:
      for i in xrange(0, len(summaryFiles)):
        print >> sys.stderr, "# Mapped reads in", summaryFiles[i], "=", mappedReads[i]
      print >> sys.stderr, "# Total mapped reads =", sum(mappedReads)
  return humanMappedReads

